import cv2
import numpy as np
import base64


def find_template_position(screenshot_base64, template_base64, threshold=0.8):
    """
    在屏幕截图中查找模板图片位置并返回中心点坐标
    
    Args:
        screenshot_base64 (str): 屏幕截图的base64编码
        template_base64 (str): 模板图片的base64编码
        threshold (float): 匹配阈值，默认为0.8
    
    Returns:
        tuple: (center_x, center_y) 中心点坐标，如果未找到则返回(None, None)
    """
    
    # 将base64编码的模板图片转换为OpenCV图像
    template_data = base64.b64decode(template_base64)
    template_np = np.frombuffer(template_data, np.uint8)
    template = cv2.imdecode(template_np, cv2.IMREAD_GRAYSCALE)
    
    # 将base64编码的屏幕截图转换为OpenCV图像
    screenshot_data = base64.b64decode(screenshot_base64)
    screenshot_np = np.frombuffer(screenshot_data, np.uint8)
    screenshot = cv2.imdecode(screenshot_np, cv2.IMREAD_GRAYSCALE)
    
    # 执行模板匹配
    result = cv2.matchTemplate(screenshot, template, cv2.TM_CCOEFF_NORMED)
    
    # 查找最佳匹配位置
    min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
    
    # 如果匹配度超过阈值，则返回中心点坐标
    if max_val >= threshold:
        # 获取模板尺寸
        template_h, template_w = template.shape
        
        # 计算中心点坐标
        center_x = max_loc[0] + template_w // 2
        center_y = max_loc[1] + template_h // 2
        
        print(f"找到匹配位置: ({center_x}, {center_y}), 匹配度: {max_val}")
        return (center_x, center_y)
    else:
        print(f"未找到匹配位置，最高匹配度: {max_val}")
        return (None, None)


def find_sift_position(screenshot_base64, template_base64, min_matches=10):
    """
    使用SIFT特征点匹配在屏幕截图中查找模板图片位置并返回中心点坐标
    
    Args:
        screenshot_base64 (str): 屏幕截图的base64编码
        template_base64 (str): 模板图片的base64编码
        min_matches (int): 最少匹配点数，默认为10
    
    Returns:
        tuple: (center_x, center_y) 中心点坐标，如果未找到则返回(None, None)
    """
    
    # 将base64编码的模板图片转换为OpenCV图像
    template_data = base64.b64decode(template_base64)
    template_np = np.frombuffer(template_data, np.uint8)
    template = cv2.imdecode(template_np, cv2.IMREAD_GRAYSCALE)
    
    # 将base64编码的屏幕截图转换为OpenCV图像
    screenshot_data = base64.b64decode(screenshot_base64)
    screenshot_np = np.frombuffer(screenshot_data, np.uint8)
    screenshot = cv2.imdecode(screenshot_np, cv2.IMREAD_GRAYSCALE)
    
    # 创建SIFT对象
    sift = cv2.SIFT_create()
    
    # 检测关键点并计算描述符
    kp1, des1 = sift.detectAndCompute(template, None)
    kp2, des2 = sift.detectAndCompute(screenshot, None)
    
    if des1 is None or des2 is None:
        print("未能检测到足够的特征点")
        return (None, None)
    
    # 使用FLANN匹配器进行特征匹配
    FLANN_INDEX_KDTREE = 1
    index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
    search_params = dict(checks=50)
    flann = cv2.FlannBasedMatcher(index_params, search_params)
    
    matches = flann.knnMatch(des1, des2, k=2)
    
    # 应用 Lowe's ratio test 来过滤好的匹配点
    good_matches = []
    for pair in matches:
        if len(pair) == 2:
            m, n = pair
            if m.distance < 0.7 * n.distance:
                good_matches.append(m)
    
    # 如果有足够的匹配点，则计算位置
    if len(good_matches) > min_matches:
        # 提取匹配点的坐标
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)
        
        # 计算单应性矩阵
        M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
        
        if M is not None:
            # 获取模板图片的尺寸
            h, w = template.shape
            
            # 定义模板图片的四个角点
            pts = np.float32([[0, 0], [0, h-1], [w-1, h-1], [w-1, 0]]).reshape(-1, 1, 2)
            
            # 使用单应性矩阵透视变换角点
            dst = cv2.perspectiveTransform(pts, M)
            
            # 计算中心点坐标
            center_x = int(np.mean(dst[:, 0, 0]))
            center_y = int(np.mean(dst[:, 0, 1]))
            
            print(f"找到匹配位置: ({center_x}, {center_y}), 匹配点数: {len(good_matches)}")
            return (center_x, center_y)
    
    print(f"未找到足够匹配点，匹配点数: {len(good_matches)}")
    return (None, None)


# 示例用法
if __name__ == "__main__":
    # 读取屏幕截图和模板图片并转为base64编码
    with open(r"D:\PycharmProjects\auto-script\screentshot\at.png", "rb") as f:
        screenshot_base64 = base64.b64encode(f.read()).decode('utf-8')
    
    with open(r"D:\PycharmProjects\auto-script\template\at.png", "rb") as f:
        template_base64 = base64.b64encode(f.read()).decode('utf-8')
        print("模板图片转为base64编码:")
        print(template_base64)
    
    # 使用模板匹配查找位置
    center_x, center_y = find_template_position(screenshot_base64, template_base64)
    
    if center_x is None or center_y is None:
        print("模板匹配失败，尝试使用SIFT特征匹配")
        center_x, center_y = find_sift_position(screenshot_base64, template_base64)
        
    if center_x is not None and center_y is not None:
        print(f"成功找到目标位置: ({center_x}, {center_y})")
        # 在此处添加点击操作或其他处理
    else:
        print("未能找到目标位置")